import logging
import lzma
from io import BytesIO
from typing import Union

from ofrak.component.packer import Packer
from ofrak.component.unpacker import Unpacker
from ofrak.resource import Resource
from ofrak.core.binary import GenericBinary
from ofrak.core.magic import MagicMimePattern, MagicDescriptionPattern
from ofrak_type.range import Range

LOGGER = logging.getLogger(__name__)


class LzmaData(GenericBinary):
    """
    An lzma binary blob.
    """

    async def get_child(self) -> GenericBinary:
        return await self.resource.get_only_child_as_view(GenericBinary)


class XzData(GenericBinary):
    """
    An xz binary blob.
    """

    async def get_child(self) -> GenericBinary:
        return await self.resource.get_only_child_as_view(GenericBinary)


class LzmaUnpacker(Unpacker[None]):
    """
    Decompresses LZMA (Lempel-Ziv-Markov chain Algorithm) or XZ compressed data. Use for .lzma or
    .xz files, or when encountering LZMA-compressed sections in firmware images. Common in router
    firmware and embedded Linux distributions.
    """

    id = b"LzmaUnpacker"
    targets = (LzmaData, XzData)
    children = (GenericBinary,)

    async def unpack(self, resource: Resource, config=None):
        file_data = BytesIO(await resource.get_data())

        format = lzma.FORMAT_AUTO

        if resource.has_tag(XzData):
            format = lzma.FORMAT_XZ
        elif resource.has_tag(LzmaData):
            format = lzma.FORMAT_ALONE

        lzma_entry_data = None
        compressed_data = file_data.read()

        try:
            lzma_entry_data = lzma.decompress(compressed_data, format)
        except lzma.LZMAError:
            LOGGER.info("Initial LZMA decompression failed. Trying with null bytes stripped")
            lzma_entry_data = lzma.decompress(compressed_data.rstrip(b"\x00"), format)

        if lzma_entry_data is not None:
            await resource.create_child(
                tags=(GenericBinary,),
                data=lzma_entry_data,
            )
        else:
            raise lzma.LZMAError("Decompressed LZMA data is null")


class LzmaPacker(Packer[None]):
    """
    Compresses data using LZMA (Lempel-Ziv-Markov chain Algorithm) or XZ compression. Use after
    modifying decompressed LZMA/XZ data to recreate compressed files or sections. Common in
    firmware, embedded systems, and software distributions.
    """

    targets = (LzmaData, XzData)

    async def pack(self, resource: Resource, config=None):
        lzma_format, tag = await self._get_lzma_format_and_tag(resource)
        lzma_file: Union[XzData, LzmaData] = await resource.view_as(tag)

        lzma_child = await lzma_file.get_child()
        lzma_compressed = lzma.compress(await lzma_child.resource.get_data(), lzma_format)

        original_size = await lzma_file.resource.get_data_length()
        resource.queue_patch(Range(0, original_size), lzma_compressed)

    async def _get_lzma_format_and_tag(self, resource):
        if resource.has_tag(XzData):
            tag = XzData
            lzma_format = lzma.FORMAT_XZ
        elif resource.has_tag(LzmaData):
            tag = LzmaData
            lzma_format = lzma.FORMAT_ALONE
        else:
            raise TypeError(
                f"Expected target of {self.get_id().decode()} to be either XzFile or LzmaFile"
            )
        return lzma_format, tag


MagicMimePattern.register(LzmaData, "application/x-lzma")
MagicMimePattern.register(XzData, "application/x-xz")
MagicDescriptionPattern.register(LzmaData, lambda s: s.startswith("LZMA compressed data"))
MagicDescriptionPattern.register(XzData, lambda s: s.startswith("XZ compressed data"))
